package main.java.com.itlin.datastructure.avltree;

public class AVLTree {
    AVLNode root;

    //创建节点类
    static class AVLNode{
        int key;
        Object value;
        AVLNode left;
        AVLNode right;

        int height = 1;//初始高度

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right, int height) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.height = height;
        }
    }

    //计算节点的高度
    public int height(AVLNode node){
        return node == null ? 0: node.height;

    }

    //更新节点高度（新增、删除、旋转）
    private void updateHeight(AVLNode node){
        node.height = Integer.max(height(node.left), height(node.right)) + 1;
    }

    //求左右子树的高度差
    //**平衡因子**：（balance factor） 左子树高度 - 右子树高度
    private int bf(AVLNode node){
        return height(node.left) - height(node.right);
        /*
         * 平衡因子：0, 1, -1 是平衡的
         *  >1 || < -1 不平衡
         */
    }

    //右旋
    //参数： 要旋转的节点， 返回值： 旋转后新的根节点
    private AVLNode rightRotate(AVLNode node){
        //取出根节点的左节点
        AVLNode left = node.left;
        //取出根节点左节点的右节点
        AVLNode Lright = left.right;
        //将根节点的左节点设置为新的根节点
        left.right = node;
        //将原来左节点的右节点设置为根节点的左节点
        node.left = Lright;
        //先更新旋转后 高度较低的
        updateHeight(node);
        updateHeight(left);

        return left;
    }

    //左旋
    private AVLNode leftRotate(AVLNode node){
        AVLNode right = node.right;
        AVLNode Rleft = right.left;
        right.left = node;
        node.right = Rleft;
        //先更新旋转后 高度较低的
        updateHeight(node);
        updateHeight(right);
        return right;
    }

    // 先左旋左子树， 再右旋根节点
    private AVLNode leftRightRotate(AVLNode node){
        node.left = leftRotate(node.left);
        return rightRotate(node);
    }

    // 先右旋右子树， 再左旋根节点
    private AVLNode rightLeftRotate(AVLNode node){
        node.right = rightRotate(node.right);
        return leftRotate(node);
    }

    // 检查节点是否失衡， 重新平衡代码
    private  AVLNode balance(AVLNode node){
        if(node == null){
            return null;
        }
        int bf = bf(node);
        if(bf > 1 && bf(node.left) >= 0){//LL
            return rightRotate(node);
        }
        else if(bf > 1 && bf(node.left) < 0){//LR
            return leftRightRotate(node);
        }
        else if(bf < -1 && bf(node.right) > 0){//RL
            return rightLeftRotate(node);
        }else if (bf < -1 && bf(node.right) <= 0){//RR
            return leftRotate(node);
        }
        return node;
    }

    //插入节点
    public void put(int key,Object value){
        root = doPut(root,key,value);
    }

    private AVLNode doPut(AVLNode node,int key, Object value){
        // 1.找到空位，创建新节点
        if(node == null){
            return new AVLNode(key,value);
        }
        //2.key 已经存在， 更新
        if(key == node.key){
            node.value = value;
        }
        //3. 继续查找
        if(key < node.key){
            node.left = doPut(node.left, key, value);

        }else{
            node.right = doPut(node.right, key, value);
        }
        //更新高度
        updateHeight(node);
        //检查是否失衡， 重新平衡
        return balance(node);
    }

    public void remove(int key){
        root = doRemove(root,key);
    }
    private AVLNode doRemove(AVLNode node, int  key){
        //1. node == null
        if(node == null) return null;
        //2. 没找到key
        if(key <  node.key){
            node.left = doRemove(node.left , key);
        }else if(key > node.key){
            node.right = doRemove(node.right, key);
        }else{
            //3. 找到key
            // 3.1 当前节点没有孩子（子节点）
            if(node.left == null && node.right == null){
                return null;
            }
            // 3.2 右子树为空，或右子树为空 直接删除
            else if(node.left == null){
                node = node.right;
                //不立刻返回，因为要调整高度
            }else if(node.right == null){
                node = node.left;
            }
            // 3.3 左右子树都有节点， 处理后继节点的后事后 将后继节点， 替换当前节点
            else{
                AVLNode s = node.right;
                while(s.left != null){
                    s = s.left;
                }
                //s 后继节点
                s.right = doRemove(node.right, s.key);
                s.left = node.left;
                node = s;
            }

        }


        //4. 更新高度
        updateHeight(node);
        //5. balance
        return balance(node);
    }
}
